package com.campus.counseling.service.impl;

import com.campus.counseling.config.LSTMConfig;
import com.campus.counseling.service.Word2VecService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.models.embeddings.loader.WordVectorSerializer;
import org.deeplearning4j.models.word2vec.Word2Vec;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.INDArrayIndex;
import org.nd4j.linalg.indexing.NDArrayIndex;
import org.springframework.stereotype.Service;
import com.huaban.analysis.jieba.JiebaSegmenter;

import javax.annotation.PostConstruct;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.stream.Collectors;
import java.io.InputStreamReader;
import java.nio.file.Files;
import java.nio.file.Path;

@Slf4j
@Service
@RequiredArgsConstructor
public class Word2VecServiceImpl implements Word2VecService {

    private final LSTMConfig lstmConfig;
    private final Word2VecTrainerImpl word2VecTrainer;
    private Word2Vec word2Vec;
    private JiebaSegmenter segmenter;
    private static final String MODEL_PATH = "models/word2vec.bin";

    @PostConstruct
    public void init() {
        try {
            segmenter = new JiebaSegmenter();
            log.info("结巴分词器初始化成功");
            
            loadModel();
        } catch (Exception e) {
            log.error("初始化失败", e);
            throw new RuntimeException("初始化失败: " + e.getMessage());
        }
    }

    @Override
    public void trainModel(List<String> sentences) {
        try {
            log.info("开始训练Word2Vec模型...");
            word2Vec = word2VecTrainer.trainWord2Vec(sentences);
            saveModel();
            log.info("Word2Vec模型训练完成");
        } catch (Exception e) {
            log.error("Word2Vec模型训练失败", e);
            throw new RuntimeException("Word2Vec模型训练失败", e);
        }
    }

    @Override
    public void trainFromFile(String filePath) {
        try {
            File file = new File(filePath);
            if (!file.exists()) {
                throw new RuntimeException("路径不存在: " + filePath);
            }

            // 如果是目录，则处理目录下的所有txt文件
            if (file.isDirectory()) {
                log.info("检测到目录，将处理目录下所有txt文件: {}", filePath);
                processDirectory(file);
                return;
            }

            // 如果是文件，确保是txt文件
            if (!file.getName().toLowerCase().endsWith(".txt")) {
                throw new RuntimeException("不支持的文件类型，仅支持txt文件: " + filePath);
            }

            // 1. 如果存在模型则加载，不存在则创建新模型
            if (word2Vec == null) {
                File modelFile = new File(MODEL_PATH);
                if (modelFile.exists()) {
                    log.info("加载现有模型进行增量训练");
                    word2Vec = WordVectorSerializer.readWord2VecModel(modelFile);
                }
            }

            // 2. 读取文件内容
            List<String> lines = new ArrayList<>();
            try (BufferedReader reader = new BufferedReader(
                    new InputStreamReader(Files.newInputStream(file.toPath()), StandardCharsets.UTF_8))) {
                String line;
                int lineCount = 0;
                while ((line = reader.readLine()) != null) {
                    if (!line.trim().isEmpty()) {
                        lines.add(line);
                    }
                    lineCount++;
                    if (lineCount % 10000 == 0) {
                        log.info("已读取 {} 行", lineCount);
                    }
                }
            }
            log.info("成功读取训练文件，共 {} 行有效数据", lines.size());

            // 3. 训练模型
            if (!lines.isEmpty()) {
                trainModel(lines);
            }

        } catch (Exception e) {
            log.error("从文件训练Word2Vec模型失败: {}", e.getMessage());
            throw new RuntimeException("从文件训练Word2Vec模型失败: " + e.getMessage());
        }
    }

    private void processDirectory(File dir) {
        try {
            File[] files = dir.listFiles((d, name) -> name.toLowerCase().endsWith(".txt"));
            if (files == null || files.length == 0) {
                throw new RuntimeException("目录中没有找到txt文件: " + dir.getPath());
            }

            log.info("找到 {} 个txt文件", files.length);
            List<String> allLines = new ArrayList<>();
            
            for (File file : files) {
                try {
                    log.info("正在处理文件: {}", file.getName());
                    try (BufferedReader reader = new BufferedReader(
                            new InputStreamReader(Files.newInputStream(file.toPath()), StandardCharsets.UTF_8))) {
                        String line;
                        while ((line = reader.readLine()) != null) {
                            if (!line.trim().isEmpty()) {
                                allLines.add(line);
                            }
                        }
                    }
                    log.info("文件 {} 处理完成", file.getName());
                } catch (Exception e) {
                    log.error("处理文件失败: {}", file.getName(), e);
                }
            }

            log.info("目录处理完成，共读取 {} 行有效数据", allLines.size());
            if (!allLines.isEmpty()) {
                trainModel(allLines);
            }

        } catch (Exception e) {
            log.error("处理目录失败: {}", dir.getPath(), e);
            throw new RuntimeException("处理目录失败: " + e.getMessage());
        }
    }

    @Override
    public Word2Vec getModel() {
        if (word2Vec == null) {
            try {
                File modelFile = new File(MODEL_PATH);
                if (modelFile.exists()) {
                    word2Vec = WordVectorSerializer.readWord2VecModel(modelFile);
                    log.info("已加载现有Word2Vec模型");
                } else {
                    throw new RuntimeException("Word2Vec模型文件不存在");
                }
            } catch (Exception e) {
                log.error("加载Word2Vec模型失败:", e);
                throw new RuntimeException("加载Word2Vec模型失败", e);
            }
        }
        return word2Vec;
    }

    @Override
    public double[] getWordVector(String word) {
        try {
            // 如果word2Vec为空或者word2Vec中没有该单词，则返回零向量
            if (word2Vec == null || !word2Vec.hasWord(word)) {
                return new double[word2Vec.getLayerSize()];
            }
            return word2Vec.getWordVector(word);
        } catch (Exception e) {
            log.error("获取词向量失败: {}", word, e);
            return new double[word2Vec.getLayerSize()];
        }
    }

    @Override
    public boolean hasWord(String word) {
        return getModel().hasWord(word);
    }

    @Override
    public double similarity(String word1, String word2) {
        return getModel().similarity(word1, word2);
    }

    @Override
    public INDArray getSentenceVector(String sentence) {
        try {
            if (sentence == null || sentence.trim().isEmpty()) {
                return Nd4j.zeros(lstmConfig.getVectorSize());
            }

            List<String> tokens = tokenize(sentence);
            if (tokens.isEmpty()) {
                return Nd4j.zeros(lstmConfig.getVectorSize());
            }

            // 收集词向量
            List<INDArray> vectors = new ArrayList<>();
            for (String token : tokens) {
                if (word2Vec.hasWord(token)) {
                    double[] vector = word2Vec.getWordVector(token);
                    if (vector.length == lstmConfig.getVectorSize()) {
                        vectors.add(Nd4j.create(vector));
                    } else {
                        log.warn("跳过维度不匹配的词向量: word='{}', dim={}", token, vector.length);
                    }
                }
            }

            if (vectors.isEmpty()) {
                log.warn("句子中没有找到任何有效词向量: {}", sentence);
                return Nd4j.zeros(lstmConfig.getVectorSize());
            }

            // 计算平均向量
            INDArray sum = Nd4j.zeros(lstmConfig.getVectorSize());
            for (INDArray vec : vectors) {
                sum.addi(vec);
            }
            
            return sum.divi(vectors.size());

        } catch (Exception e) {
            log.error("获取句子向量失败: {}", sentence, e);
            return Nd4j.zeros(lstmConfig.getVectorSize());
        }
    }
    
    private List<String> tokenize(String text) {
        try {
            if (segmenter == null) {
                segmenter = new JiebaSegmenter();
                log.info("重新初始化结巴分词器");
            }
            
            // 使用结巴分词
            return segmenter.sentenceProcess(text.trim())
                .stream()
                .filter(token -> !token.trim().isEmpty())
                .collect(Collectors.toList());
        } catch (Exception e) {
            log.error("分词失败: {}", text, e);
            return new ArrayList<>();
        }
    }

    @Override
    public void saveModel() {
        try {
            File modelDir = new File("models");
            if (!modelDir.exists()) {
                modelDir.mkdirs();
            }
            
            File modelFile = new File(lstmConfig.getWord2vecPath());
            WordVectorSerializer.writeWord2VecModel(word2Vec, modelFile);
            log.info("Word2Vec模型保存成功: {}", modelFile.getAbsolutePath());
        } catch (Exception e) {
            log.error("Word2Vec模型保存失败", e);
            throw new RuntimeException("Word2Vec模型保存失败", e);
        }
    }

    @Override
    public void loadModel() {
        try {
            String modelPath = lstmConfig.getWord2vecPath();
            log.info("开始加载Word2Vec模型: {}", modelPath);
            
            File modelFile = new File(modelPath);
            if (!modelFile.exists()) {
                throw new RuntimeException("Word2Vec模型文件不存在: " + modelPath);
            }
            
            // 使用新的API
            word2Vec = WordVectorSerializer.readWord2VecModel(modelFile);
            log.info("Word2Vec模型加载成功");
            
        } catch (Exception e) {
            log.error("Word2Vec模型加载失败", e);
            throw new RuntimeException("Word2Vec模型加载失败: " + e.getMessage());
        }
    }
} 